/**
 * Copyright (C) 2009+ Spark Zheng
 *
 * @file        test/mc_press.c 
 * @date        2010-09-15
 * @brief	test_module
 *		mc protocol press-testing
 *		Two Method:
 *		-1. create pretty many threads
 *		-2. use event-and-nonblock client
 * @see         dserver
 * @history
 */

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <sys/types.h>
#include <sys/time.h>
#include <sys/resource.h>

#include <dlib/global.h>
#include <dlib/log.h>
#include <dlib/str.h>
#include <dlib/dtime.h>

#include "libdserver.h"

/* 1000 key-set */
#define TEST_LOOP	1000
/* 20k buffer */
#define TEST_BUF_LEN	20480
/* 1024 buffer */
#define FILE_BUF_LEN	1024

static void _handle_input(int fd, short event, void *arg);
static void _handle_output(int fd, short event, void *arg);
static void *_press_thread(void *arg);
static int press_threadpool(int concurrence);
static int press_eventpool(int concurrence);

static struct press_addr {
	const char *ip;
	in_port_t port;
} addr;
static char *key_file = NULL;

/* global test str */
static int test_str_len = 0;
static char *test_str[TEST_LOOP];
static struct timeval start, end, result;

int main(int argc, char **argv)
{
	int i;
	int line_nr = 0;
	char file_buf[FILE_BUF_LEN] = {0};

	if (argc < 3) {
		fprintf(stderr, "Usage: mc_press type concurrence test_file ip port\n");
		fprintf(stderr, "type: 1 - threadpool, 2 - eventpool\n");
		fprintf(stderr, "concurrence: max-parallel requests\n");
		fprintf(stderr, "test_file: key list file\n");
		fprintf(stderr, "ip/port: server ip & port\n");
		return -1;
	}

	int type = atoi(argv[1]);
	int concurrence = atoi(argv[2]);

	if (argc > 3) {
		key_file = argv[3];
	} else {
		key_file = "key_list.txt";
	}
	if (argc > 4) {
		addr.ip = argv[4];
	} else {
		addr.ip = "127.0.0.1";
	}
	if (argc > 5) {
		addr.port = (in_port_t)atoi(argv[5]);
	} else {
		addr.port = 10011;
	}
	
	/* read key file for TEST_LOOP data */
	FILE *fp = fopen(key_file, "r");
	if (fp == NULL) {
		fprintf(stderr, "test_file %s open error.\n", key_file);
		fprintf(stderr, "Usage: mc_press type concurrence test_file ip port\n");
		fprintf(stderr, "type: 1 - threadpool, 2 - eventpool\n");
		fprintf(stderr, "concurrence: max-parallel requests\n");
		fprintf(stderr, "test_file: key list file\n");
		fprintf(stderr, "ip/port: server ip & port\n");
		return -2;
	}

	while ((line_nr < TEST_LOOP) && (fgets(file_buf, FILE_BUF_LEN, fp) != NULL)) {
		test_str[line_nr] = strdup(str_trim(file_buf));
		test_str_len = line_nr;
		line_nr++;
		memset(file_buf, 0, FILE_BUF_LEN);
	}

	if (test_str_len == 0) {
		fprintf(stderr, "test_file %s no keys error.\n", key_file);
		fprintf(stderr, "Usage: mc_press type concurrence test_file ip port\n");
		fprintf(stderr, "type: 1 - threadpool, 2 - eventpool\n");
		fprintf(stderr, "concurrence: max-parallel requests\n");
		fprintf(stderr, "test_file: key list file\n");
		fprintf(stderr, "ip/port: server ip & port\n");
		return -3;
	}

	fprintf(stderr, "PRESS BEGIN:\b");
	if (type == 1) {
		struct rlimit rlimit;

		getrlimit(RLIMIT_NPROC, &rlimit);
		fprintf(stderr, "cur NPROC: %lu, max NPROC: %lu\n",
				rlimit.rlim_cur, rlimit.rlim_max);
		if ((uint32_t)concurrence > rlimit.rlim_cur &&
				(uint32_t)concurrence < rlimit.rlim_max) {
			rlimit.rlim_cur = concurrence + 1;
			setrlimit(RLIMIT_NPROC, &rlimit);
		}

		press_threadpool(concurrence);
	} else {
		press_eventpool(concurrence);
	}
	fprintf(stderr, "PRESS END.\n");

	/* free key_str list */
	for (i = 0; i < test_str_len; i++) {
		free(test_str[i]);
		test_str[i] = NULL;
	}

	return 0;
}

static pthread_mutex_t mutex;

static int press_threadpool(int concurrence)
{
	int i;
	pthread_t pid[concurrence];
	pthread_attr_t attr;

	pthread_attr_init(&attr);
	//pthread_attr_setstacksize(&attr, 128); /* 1024 * 1024  * 3 / 128 = 24576 */
	pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);
	pthread_mutex_init(&mutex, NULL);

	for (i = 0; i < concurrence; i++) {
		pthread_mutex_lock(&mutex);
		if (pthread_create(&pid[i], &attr, _press_thread, &i) != 0) {
			fprintf(stderr, "create press thread err[%d]:%s\n", i, strerror(errno));
			pthread_mutex_unlock(&mutex);
		} else {
			fprintf(stderr, "create press thread[%d] success.\n", i);
		}
	}

	pause();
	/*
	for (i = 0; i < concurrence; i++) {
		pthread_join(pid[i], NULL);
	}
	*/
	
	return 0;
}

struct client {
	int		sfd;
	struct event	event;
};

static int press_eventpool(int concurrence)
{
	int i;
	struct client client[concurrence];

	event_init();

	for (i = 0; i < concurrence; i++) {
		client[i].sfd = dserver_tcpclient_init(addr.ip, addr.port, 30);
		if (client[i].sfd < 0) {
			fprintf(stderr, "create tcpclient error.\n");
			continue;
		} else {
			fprintf(stderr, "create sfd = %d.\n", client[i].sfd);
		}

		event_set(&client[i].event, client[i].sfd, EV_WRITE, _handle_output, &client[i]);
		event_add(&client[i].event, NULL);
	}

	fprintf(stderr, "eventpool begin loop.\n");
	event_dispatch();

	return 0;
}

static void *_press_thread(void *arg)
{
	int i;
	int sfd;
	int ret = *(int *)arg;
	size_t wbuf_len = 0;
	char wbuf[FILE_BUF_LEN + 1] = {0};
	char rbuf[TEST_BUF_LEN + 1] = {0};

	fprintf(stderr, "IN thread %d.\n", ret);
	pthread_mutex_unlock(&mutex);

	sfd = dserver_tcpclient_init(addr.ip, addr.port, 30);
	if (sfd < 0) {
		fprintf(stderr, "create tcpclient error.\n");
		return NULL;
	} else {
		fprintf(stderr, "create sfd = %d.\n", sfd);
	}

	while (1) {
		i = random() % test_str_len;
		wbuf_len = snprintf(wbuf, FILE_BUF_LEN, "get %s\n\r", test_str[i]);
		fprintf(stderr, "tcp write data[%s].\n", wbuf);

		DLIB_UTIMER_RECORD(&start);
		ret = dserver_tcpdata_write(sfd, wbuf, wbuf_len, 1);
		if (ret < 0) {
			continue;	
		}
		ret = dserver_tcpdata_read(sfd, rbuf, TEST_BUF_LEN, 1);
		if (ret < 0) {
			continue;	
		}
		DLIB_UTIMER_RECORD(&end);
		DLIB_UTIMER_SUB(&end, &start, &result);
		fprintf(stdout, "%ld\n", result.tv_sec * 1000000 + result.tv_usec);

		rbuf[ret] = 0;
		fprintf(stderr, "tcp read data[%s].\n", rbuf);
	//	usleep(500);
	}

	dserver_tcpclient_exit(sfd);
}

static void _handle_input(int fd, short event, void *arg)
{
	int ret;
	struct client *client = (struct client *)arg;
	char buf[TEST_BUF_LEN + 1] = {0};

	UNUSED_PARAM(fd);
	UNUSED_PARAM(event);

	ret = read(client->sfd, buf, TEST_BUF_LEN);
	DLIB_UTIMER_RECORD(&end);
	DLIB_UTIMER_SUB(&end, &start, &result);
	fprintf(stdout, "%ld\n", result.tv_sec * 1000000 + result.tv_usec);

	if (ret > 0) {
		fprintf(stderr, "tcp read data[%s].\n", buf);
		event_set(&client->event, client->sfd, EV_WRITE, _handle_output, client);
		event_add(&client->event, NULL);
	} else if (ret == 0) {
		fprintf(stderr, "tcp read FIN, close.\n");
		dserver_tcpclient_exit(client->sfd);
	} else {
		event_set(&client->event, client->sfd, EV_READ, _handle_input, client);
		event_add(&client->event, NULL);
	}

	return;
}

static void _handle_output(int fd, short event, void *arg)
{
	int ret;
	struct client *client = (struct client *)arg;
	int i = random() % test_str_len;
	size_t buf_len = 0;
	char buf[FILE_BUF_LEN + 1] = {0};

	UNUSED_PARAM(fd);
	UNUSED_PARAM(event);

	buf_len = snprintf(buf, FILE_BUF_LEN, "get %s\n\r", test_str[i]);
	fprintf(stderr, "tcp write data[%s].\n", buf);

	DLIB_UTIMER_RECORD(&start);
	ret = write(client->sfd, buf, buf_len);
	if (ret > 0) {
		event_set(&client->event, client->sfd, EV_READ, _handle_input, client);
		event_add(&client->event, NULL);
	} else {
		event_set(&client->event, client->sfd, EV_WRITE, _handle_output, client);
		event_add(&client->event, NULL);
	}

	return;
}

